package com.example.audio_test.utils

import android.media.MediaCodec
import android.media.MediaExtractor
import android.media.MediaFormat
import android.media.MediaMuxer
import java.io.File
import java.nio.ByteBuffer
import kotlin.math.max

/**
 *视频合成
 *
 * @author : Kun
 * @date : 2021/6/23
 * @time : 16:47
 */
class VideoMuxer {

    private val MAX_BUFF_SIZE = 1048576

    private val mReaderBuffer = ByteBuffer.allocate(MAX_BUFF_SIZE)

    /**
     * 多个视频文件合成一个
     */
    fun muxVideoList(videoList: Array<File>, outPath: String, finish: () -> Unit) {
        // 创建MediaMuxer
        val mediaMuxer = MediaMuxer(outPath, MediaMuxer.OutputFormat.MUXER_OUTPUT_MPEG_4)
        mediaMuxer.setOrientationHint(90)

        //找到文件的音频格式和视频格式
        var findAudioFormat = false
        var findVideoFormat = false

        var audioFormat: MediaFormat? = null
        var videoFormat: MediaFormat? = null

        for (file in videoList) {
            //创建视频提取器
            val mediaExtractor = MediaExtractor()
            mediaExtractor.setDataSource(file.absolutePath)

            if (!findAudioFormat) {
                //提取音频格式
                audioFormat = findFormat(mediaExtractor, "audio/")
                findAudioFormat = audioFormat != null
            }

            if (!findVideoFormat) {
                //提取视频格式
                videoFormat = findFormat(mediaExtractor, "video/")
                findVideoFormat = videoFormat != null
            }

            //释放资源
            mediaExtractor.release()

            if (findAudioFormat && findVideoFormat) {
                break;
            }
        }

        var mediaMuxerAudioTrackIndex = 0
        if (findAudioFormat) {
            //合成文件添加指定格式音轨
            mediaMuxerAudioTrackIndex = mediaMuxer.addTrack(audioFormat!!)
        }

        var mediaMuxerVideoIndex = 0
        if (findVideoFormat) {
            //合成文件添加指定格式视频
            mediaMuxerVideoIndex = mediaMuxer.addTrack(videoFormat!!)
        }
        //开始合成
        mediaMuxer.start()

        //音视频时间戳
        var audioPts = 0L
        var videoPts = 0L

        for (file in videoList) {
            var hasAudio = false
            var hasVideo = false

            //文件音轨
            val audioMediaExtractor = MediaExtractor()
            audioMediaExtractor.setDataSource(file.absolutePath)
            val audioTrackIndex = findTrackIndex(audioMediaExtractor, "audio/")
            if (audioTrackIndex >= 0) {
                audioMediaExtractor.selectTrack(audioTrackIndex)
                hasVideo = true
            }

            //视频轨道
            val videoMediaExtractor = MediaExtractor()
            videoMediaExtractor.setDataSource(file.absolutePath)
            val videoTrackIndex = findTrackIndex(videoMediaExtractor, "video/")
            if (videoTrackIndex >= 0) {
                videoMediaExtractor.selectTrack(videoTrackIndex)
                hasVideo = true
            }

            //音视频都没有，跳过该文件
            if (!hasAudio && !hasVideo) {
                audioMediaExtractor.release()
                videoMediaExtractor.release()
                continue
            }

            //写入音频
            if (hasAudio) {
                var hasDone = false
                var lastPts = 0L
                while (!hasDone) {
                    mReaderBuffer.rewind()
                    val frameSize = audioMediaExtractor.readSampleData(mReaderBuffer, 0)
                    if (frameSize < 0) {
                        hasDone = true
                    } else {
                        val bufferInfo = MediaCodec.BufferInfo()
                        bufferInfo.offset = 0
                        bufferInfo.size = frameSize
                        bufferInfo.presentationTimeUs = audioPts + audioMediaExtractor.sampleTime

                        if (audioMediaExtractor.sampleFlags and MediaCodec.BUFFER_FLAG_KEY_FRAME != 0) {
                            bufferInfo.flags = MediaCodec.BUFFER_FLAG_KEY_FRAME
                        }

                        mReaderBuffer.rewind()
                        mediaMuxer.writeSampleData(
                            mediaMuxerAudioTrackIndex,
                            mReaderBuffer,
                            bufferInfo
                        )
                        audioMediaExtractor.advance()
                        if (audioMediaExtractor.sampleTime > 0) {
                            lastPts = audioMediaExtractor.sampleTime
                        }
                    }
                }
                audioPts += lastPts
                audioMediaExtractor.release()
            }

            if (hasVideo) {
                var hasDone = false
                var lastPts = 0L
                while (!hasDone) {
                    mReaderBuffer.rewind()
                    //读取数据
                    val frameSize = videoMediaExtractor.readSampleData(mReaderBuffer, 0)
                    if (frameSize < 0) {
                        hasDone = true
                    } else {
                        val bufferInfo = MediaCodec.BufferInfo()
                        bufferInfo.offset = 0
                        bufferInfo.size = frameSize
                        bufferInfo.presentationTimeUs = videoPts + videoMediaExtractor.sampleTime
                        if ((videoMediaExtractor.sampleFlags and MediaCodec.BUFFER_FLAG_KEY_FRAME) != 0) {
                            //包含关键帧的话
                            bufferInfo.flags = MediaCodec.BUFFER_FLAG_KEY_FRAME
                        }
                        mReaderBuffer.rewind()
                        //混合器写入数据
                        mediaMuxer.writeSampleData(
                            mediaMuxerAudioTrackIndex,
                            mReaderBuffer,
                            bufferInfo
                        )
                        //移动到下一个数据样本
                        videoMediaExtractor.advance()
                        if (videoMediaExtractor.sampleTime > 0) {
                            lastPts = videoMediaExtractor.sampleTime
                        }
                    }
                }
                videoPts += lastPts
                videoMediaExtractor.release()
            }

            //合并完一个文件 对齐时间戳
            audioPts = max(audioPts, videoPts)
            videoPts = audioPts
        }

        mediaMuxer.stop()
        mediaMuxer.release()
        finish.invoke()
    }

    /**
     * 提取对应的媒体格式
     */
    private fun findFormat(mediaExtractor: MediaExtractor, prefix: String): MediaFormat? {
        for (i in 0..mediaExtractor.trackCount) {
            val format = mediaExtractor.getTrackFormat(i)
            if (format.getString("mime")!!.startsWith(prefix)) {
                return format
            }
        }
        return null
    }

    /**
     * 查找对应轨道
     */
    private fun findTrackIndex(mediaExtractor: MediaExtractor, prefix: String): Int {
        for (i in 0..mediaExtractor.trackCount) {
            val format = mediaExtractor.getTrackFormat(i)
            if (format.getString("mime")!!.startsWith(prefix)) {
                return i
            }
        }
        return -1
    }
}